#直接调用返回当前可用内存/显存
import torch
import pynvml
import psutil

def get_avaliable_memory(device):
    if device==torch.device('cuda:0'):
        pynvml.nvmlInit()
        handle = pynvml.nvmlDeviceGetHandleByIndex(0)        # 0表示第一块显卡
        meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
        ava_mem=round(meminfo.free/1024**2)
        print('current available video memory is' +' : '+ str(round(meminfo.free/1024**2)) +' MIB')

    elif device==torch.device('cpu'):
        mem = psutil.virtual_memory()
        print('current available memory is' +' : '+ str(round(mem.used/1024**2)) +' MIB')
        ava_mem=round(mem.used/1024**2)

    elif device==torch.device('mps'):
        mem = psutil.virtual_memory()
        print('current available memory is' +' : '+ str(round(mem.used/1024**2)) +' MIB')
        ava_mem = round(mem.used / 1024 ** 2)

    return ava_mem